/*
 *  cylk.c
 *  cylk
 *
 *  Created by Benjamin Heasly on 7/6/09.
 *  Copyright 2009 benjamin.heasly@gmail.com. Seattle, WA.  All rights reserved.
 *
 */

#include "cylk.h"
#include "cylk_basicUtilities.h"
#include "cylk_regionDetection.h"

#pragma mark AFNI datasets

void runCylk(THD_3dim_dataset *inputDataSet,
			 PLUGIN_interface *plint,
			 int do_projection,
			 float fit_tolerance,
			 int do_rotation, 
			 int seed_at_max_voxel, 
			 float region_thresh_pct, 
			 int xSeed,
			 int ySeed, 
			 int zSeed) {
	
	fprintf(stderr, "---- run cylk\n");
	fprintf(stderr, "---- ---- do projection=%d\n", do_projection);
	fprintf(stderr, "---- ---- fit tolerance=%.2f\n", fit_tolerance);
	fprintf(stderr, "---- ---- do rotation=%d\n", do_rotation);
	fprintf(stderr, "---- ---- seed_at_max_voxel=%d\n", seed_at_max_voxel);
	fprintf(stderr, "---- ---- region_thresh(%%)=%d\n", region_thresh_pct);
	fprintf(stderr, "---- ---- saggital(x)=%d\n", xSeed);
	fprintf(stderr, "---- ---- coronal(y)=%d\n", ySeed);
	fprintf(stderr, "---- ---- axial(z)=%d\n", zSeed);
	
	int sizeOfRegion=0, regionFailed=1, projectionFailed=1;
	float r, phi, tilt, anchor[3];
	cylkImage dataImage, regionImage, projectionImage;
	THD_3dim_dataset *regionDataset, *projectionDataset, *rotatedDataset;
	
	// load dataset and build 3D images to work with
	DSET_load(inputDataSet);
	dataImage = new_cylkImageFromAFNIDataset(inputDataSet);
	regionImage = new_cylkImageLikeExisting(&dataImage);
	projectionImage = new_cylkImageLikeExisting(&dataImage);
	
	// locate voxel seed for growing region
	if (seed_at_max_voxel) {
		indexToSubscripts(&dataImage, &xSeed, &ySeed, &zSeed, indexOfMaxImageVoxel(&dataImage));
		fprintf(stderr, "---- ---- maximum voxel at x=%d, y=%d, z=%d\n", xSeed, ySeed, zSeed);
	} else if(!AFNI_yesenv("AFNI_LEFT_IS_LEFT")) {
		fprintf(stderr, "---- ---- flip y %d->%d\n", ySeed, dataImage.ny - ySeed);
		ySeed = dataImage.ny - ySeed;
	}
	
	// grow region near voxel seed, save to new dataset
	sizeOfRegion = growRegionWithThresholdFromSeedXYZ(&dataImage, 
													  &regionImage, 
													  region_thresh_pct, 
													  xSeed, 
													  ySeed, 
													  zSeed);
	regionFailed = sizeOfRegion <= 0;
	if (!regionFailed) {
		regionDataset = copyDatasetWithNameReplaceBrickWithImage(inputDataSet, 
																 plint, 
																 "cylk_region",
																 &regionImage);
	}
	
	// project cylinder from region, save to new dataset
	if (do_projection && !regionFailed) {
		projectionFailed = projectCylinderWithFitTolerance(&regionImage, 
														   &projectionImage, 
														   fit_tolerance, 
														   &r,
														   &phi,
														   &tilt,
														   anchor);
	}
	if (!projectionFailed) {
		projectionDataset = copyDatasetWithNameReplaceBrickWithImage(inputDataSet, 
																	 plint, 
																	 "cylk_projection",
																	 &projectionImage);
	}
	
	// rotate all datasets in line with projected cylinder, save 3 rotated datasets
	if (do_rotation && !projectionFailed) {
		float rotateABC[3], shiftXYZ[3];
		rotateABC[0] = 0;
		rotateABC[1] = -tilt;
		rotateABC[2] = -phi;
		shiftXYZ[0] = 0;
		shiftXYZ[1] = 0;
		shiftXYZ[2] = 0;
		rotateImageVolume(&dataImage, rotateABC, shiftXYZ);
		rotatedDataset = copyDatasetWithNameReplaceBrickWithImage(inputDataSet, 
																  plint, 
																  "cylk_rotation",
																  &dataImage);
		
		rotateImageVolume(&regionImage, rotateABC, shiftXYZ);
		rotatedDataset = copyDatasetWithNameReplaceBrickWithImage(inputDataSet, 
																  plint, 
																  "cylk_rotation_region",
																  &regionImage);
		
		rotateImageVolume(&projectionImage, rotateABC, shiftXYZ);
		rotatedDataset = copyDatasetWithNameReplaceBrickWithImage(inputDataSet, 
																  plint, 
																  "cylk_rotation_projection",
																  &projectionImage);				
	}
	destroy_cylkImage(&dataImage);
	destroy_cylkImage(&regionImage);
	destroy_cylkImage(&projectionImage);
}

THD_3dim_dataset* copyDatasetWithNameReplaceBrickWithImage(THD_3dim_dataset *inputDataSet, 
														   PLUGIN_interface *plint,
														   char* name, 
														   cylkImage *im) {
	fprintf(stderr, "---- copy dataset to '%s'\n", name);
	
	void *afniData;
	int nBytes;
	
	THD_3dim_dataset *newDataSet;
	newDataSet = EDIT_full_copy(inputDataSet, name);
	if(NULL!=im) {
		// AFNI should own this memory
		nBytes = im->nVoxels*sizeof(*(im->pt));
		afniData = malloc(nBytes);
		memcpy(afniData, im->pt, nBytes);
		EDIT_substitute_brick(newDataSet, 0, MRI_short, afniData);
	}
	PLUTO_add_dset(plint, newDataSet, DSET_ACTION_NONE);
	return(newDataSet);
}

void rotateImageVolume(cylkImage *im, 
					   float *rotateABC, 
					   float *shiftXYZ) {
	fprintf(stderr, "---- rotate image by [%.3f %.3f %.3f] and shift by [%.3f %.3f %.3f]\n", 
			rotateABC[0], rotateABC[1], rotateABC[2], shiftXYZ[0], shiftXYZ[1], shiftXYZ[2]);
	
	float *workingFloats;
	
	// convert short to float for rotation
	workingFloats = calloc(im->nVoxels, sizeof(*workingFloats));
	EDIT_coerce_type(im->nVoxels,
					 MRI_short, 
					 im->pt,
					 MRI_float,
					 workingFloats);
	
	// do rotation
	THD_rota_vol(im->nx, im->ny, im->nz,
				 im->dx, im->dy, im->dz,
				 workingFloats,
				 2, rotateABC[2], 1, rotateABC[1], 0, rotateABC[0],
				 DELTA_BEFORE,
				 shiftXYZ[0], shiftXYZ[1], shiftXYZ[2]);
	
	// convert back to original type
	EDIT_coerce_type(im->nVoxels,
					 MRI_float,
					 workingFloats,
					 MRI_short, 
					 im->pt);
	free(workingFloats);
}

cylkImage new_cylkImageFromAFNIDataset(THD_3dim_dataset *dSet) {
    cylkImage im = new_cylkImage(dSet->daxes->nxx, 
								 dSet->daxes->nyy, 
								 dSet->daxes->nzz,
								 fabs(dSet->daxes->xxdel), 
								 fabs(dSet->daxes->yydel), 
								 fabs(dSet->daxes->zzdel));
	
	// convert or copy data from brick to this image 
	EDIT_coerce_type(im.nVoxels, DSET_BRICK_TYPE(dSet,0), DSET_ARRAY(dSet,0), MRI_short, im.pt);
    return(im);
}

#pragma mark cylinder calculations

int projectCylinderWithFitTolerance(cylkImage *dataIm,
									cylkImage *projectionIm, 
									float tolerance, 
									float *r,
									float *phi,
									float *tilt, 
									float *anchor) {
	fprintf(stderr, "---- fit cylinder with fit tolerance %.2f\n", tolerance);
	
	int status, i, bestDim, mostSlices=0, nGoodSlices[3];
	
	for(i=0; i<3; i++){
		nGoodSlices[i] = countGoodEllipseFitsForImageDimensionWithTolerance(dataIm, i, tolerance);
        if (nGoodSlices[i] > mostSlices) {
            bestDim = i;
            mostSlices = nGoodSlices[i];
        }
	}
	
    if (mostSlices >= 3)
        status = buildCylinderWithDimensionAndFitTolerance(dataIm, 
														   projectionIm, 
														   bestDim, 
														   tolerance,
														   r, 
														   phi, 
														   tilt, 
														   anchor);
    else
		status = 1;
	
	fprintf(stderr, "---- ---- cylinder fits: r=%.3f phi=%.3f, tilt=%.3f\n", r, phi, tilt);
	fprintf(stderr, "---- ---- cylinder fits: x=%.3f, y=%.3f, z=%.3f\n", anchor[0], anchor[1], anchor[2]);
    return (status);
}

int countGoodEllipseFitsForImageDimensionWithTolerance(cylkImage *im, int dim, float tolerance) {
	
	fprintf(stderr, "---- count slices along dimension %d meeting tolerance %.2f\n", dim, tolerance);
	
	int i=0, nSlices, nGoodSlices=0;
	float fitError;
	cylkSlice sl;
	
	
	switch (dim) {
		case 0:
			nSlices = im->nx;
			break;
		case 1:
			nSlices = im->ny;
			break;
		case 2:
			nSlices = im->nz;
			break;
	}
	
	for (i=0; i<nSlices; i++) {
		sl = new_cylkSliceFromCylkImage(im, i, dim);
		fitError = errorOfEllipseFitToSlice(&sl, NULL, NULL, NULL, NULL, NULL, NULL);
		nGoodSlices += (fitError <= tolerance);
		destroy_cylkSlice(&sl);
	}
	fprintf(stderr, "---- ---- counted %d good slices\n", nGoodSlices);
	return(nGoodSlices);
}

float errorOfEllipseFitToSlice(cylkSlice *sl,
							   float *r, 
							   float *phi, 
							   float *tilt, 
							   float *centX, 
							   float *centY, 
							   float *ecc) {
	
	int i;
	float *rx, *ry;
	float x_bar, y_bar, phix, ltilt, af, bf, a, b;
	float u11, u20, u02, u20p, u02p;
	float Denom, A11, A12, A21, A22;
	float error, diff=0, area=0;
	cylkSlice fitEllipseSlice;
	
	//build rx and ry
	rx = malloc(sl->nVoxels * sizeof(*rx));
	ry = malloc(sl->nVoxels * sizeof(*ry));
	for (i=0; i<sl->nVoxels; i++) {
		if (sl->pt[i]) {
			rx[i] = (i % sl->nx) * sl->dx;
			ry[i] = (i / sl->nx) * sl->dy;
		} else {
			rx[i] = NAN;
			ry[i] = NAN;
		}
	}
	x_bar = meanOfValues(rx, sl->nVoxels);
	y_bar = meanOfValues(ry, sl->nVoxels);
	
	// subtract the center
	for(i=0; i<sl->nVoxels; i++) {
		if(!isnan(rx[i])) {
			rx[i] -= x_bar;
			ry[i] -= y_bar;
		}
	}
	
	u11 = rawMomentOfXToPAndYToQ(rx, 1, sl->dx, ry, 1, sl->dy, sl->nVoxels);
	u20 = rawMomentOfXToPAndYToQ(rx, 2, sl->dx, ry, 0, sl->dy, sl->nVoxels);
	u02 = rawMomentOfXToPAndYToQ(rx, 0, sl->dx, ry, 2, sl->dy, sl->nVoxels);
	
	// get rotation of ellipse in slice
	// rotation is only valid [-90,90]
	// all other angles are degenerate to that range
	if (fabs (u11) < 1e-9)
		phix = 0.0;
	else
		phix = atan (2.0 * u11 / (u20 - u02)) / 2.0;
	
	if (u02 > u20)
		phix = phix - M_PI_2;
	
	if (phix < (-M_PI_2))
		phix = phix + M_PI;
	
	Denom = 1 - 2*pow(sin(phix), 2);
	A11 = -pow(cos(phix), 2) / Denom;
	A12 = pow(sin(phix), 2) / Denom;
	A21 = pow(sin(phix), 2) / Denom;
	A22 = -pow(cos(phix), 2) / Denom;
	u20p = A11 * u02 + A12 * u20;
	u02p = A21 * u02 + A22 * u20;
	
	af = pow(4.0/M_PI, 0.25) * pow(pow(u20p, 3)/u02p, 0.125);
	bf = pow(4.0/M_PI, 0.25) * pow(pow(u02p, 3)/u20p, 0.125);
	
	a = MAX(af, bf);
	b = MIN(af, bf);
	ltilt = fabs(acos(b/a));
	fitEllipseSlice = new_cylkSliceLikeSliceWithFilledEllipse(sl, a, b, phix, x_bar, y_bar);
	
	for (i=0; i<sl->nVoxels; i++) {
		if (sl->pt[i] || fitEllipseSlice.pt[i]) {
			diff += abs(sl->pt[i] - fitEllipseSlice.pt[i]);
			if (sl->pt[i])
				area += sl->pt[i];
		}
	}
	
	destroy_cylkSlice(&fitEllipseSlice);
	free (rx);
	free (ry);
	
	// assign fit parameters ?
	if(NULL != r) *r = b;
	if(NULL != phi) *phi = phix;
	if(NULL != tilt) *tilt = ltilt;
	if(NULL != centX) *centX = x_bar;
	if(NULL != centY) *centY = y_bar;
	if(NULL != ecc) *ecc = sqrt(1-pow(b/a, 2));
	
	// return error
	error = diff/area;
	return (error);
}

int buildCylinderWithDimensionAndFitTolerance(cylkImage *dataIm,
											  cylkImage *projectIm, 
											  int dim, 
											  float tolerance, 
											  float *rOut, 
											  float *phiOut, 
											  float *tiltOut, 
											  float *anchor) {
	fprintf(stderr, "---- build cylinder using dimension %d and tolerance %.2f\n", dim, tolerance);
	
	int i, nSlices;
	int *sliceIsGood;
	int sign_x, sign_y, count;
	float *r, *phi, *tilt, *centX, *centY, *ecc, *zLocs;
	float fitError, sliceThickness;
	float tx, ty;
	float xLine[2], yLine[2];
	float centerMx, centerMy, centerBx, centerBy;
	float centerPhi, centerTilt;
	float cylMx, cylMy, cylBx, cylBy;
	float cylPhi, cylTilt;
	float slopeX, interceptX, slopeY, interceptY;
	float xCenter, yCenter, zCenter, majorAxisLength;
	float axis[3], axisNorm;
	cylkSlice sl, ellipseSlice;
	
	switch (dim) {
		case 0:
			nSlices = dataIm->nx;
			sliceThickness = dataIm->dx;
			break;
		case 1:
			nSlices = dataIm->ny;
			sliceThickness = dataIm->dy;
			break;
		case 2:
			nSlices = dataIm->nz;
			sliceThickness = dataIm->dz;
			break;
	}
	
	sliceIsGood = calloc(nSlices, sizeof(*sliceIsGood));
	r = calloc(nSlices, sizeof(*r));
	phi = calloc(nSlices, sizeof(*phi));
	tilt = calloc(nSlices, sizeof(*tilt));
	centX = calloc(nSlices, sizeof(*centX));
	centY = calloc(nSlices, sizeof(*centY));
	ecc = calloc(nSlices, sizeof(*ecc));
	zLocs = calloc(nSlices, sizeof(*zLocs));	
	
	for (i=0; i<nSlices; i++) {
		sl = new_cylkSliceFromCylkImage(dataIm, i, dim);
		fitError = errorOfEllipseFitToSlice(&sl, &r[i], &phi[i], &tilt[i], &centX[i], &centY[i], &ecc[i]);
		sliceIsGood[i] = fitError < tolerance;
		destroy_cylkSlice(&sl);
		zLocs[i] = i * sliceThickness;
	}
	
	// map phi into the correct quadrant
	count=0;
	for (i=0; i<nSlices; i++) {
		if (sliceIsGood[i]) {
			if (count==1) {
				if ((centX[i]-tx) >= 0)
					sign_x = 1;
				else 
					sign_x = -1;
				
				if ((centY[i]-ty) >= 0) 
					sign_y = 1;
				else 
					sign_y = -1;
				
				break;
			}
			count++;
			tx = centX[i];
			ty = centY[i];
		}
	}
	if (sign_x == -1) cylPhi = M_PI + sign_x * sign_y * meanOfAbsoluteValues(phi, nSlices);
	else cylPhi = 2 * M_PI + sign_x * sign_y * meanOfAbsoluteValues(phi, nSlices);
	cylPhi = cylPhi - (2 * M_PI) * floor (cylPhi / 2 / M_PI);
	
	*rOut   = meanOfValues (r, nSlices);
	cylTilt = meanOfValues (tilt, nSlices);
	
	cylMx = sliceThickness * cos (cylPhi) * fabs (tan (cylTilt));
	cylMy = sliceThickness * sin (cylPhi) * fabs (tan (cylTilt));
	cylBx = yInterceptOfXYPointsGivenSlope (centX, zLocs, sliceIsGood, cylMx, nSlices);
	cylBy = yInterceptOfXYPointsGivenSlope (centY, zLocs, sliceIsGood, cylMy, nSlices);
	
	fitLineToXYPoints (centX, zLocs, sliceIsGood, nSlices, xLine);
	fitLineToXYPoints (centY, zLocs, sliceIsGood, nSlices, yLine);
	centerMx = xLine[0];
	centerMy = yLine[0];
	centerBx = xLine[1];
	centerBy = yLine[1];
	
	centerPhi = atan2f (centerMy,centerMx);
	centerTilt = atan2f (sqrt (pow (centerMy, 2) + pow (centerMx, 2)), sliceThickness);
	centerPhi = (centerPhi > 0) ? centerPhi : centerPhi + 2 * M_PI;
	
	// free memory
	free (r);
	free (phi);
	free (tilt);
	free (centX);
	free (centY);
	free (ecc);
	free (sliceIsGood);
	
	if (1) {
		// Use phi and tilt from the elipses but center line from the centers
		slopeX = centerMx;
		slopeY = centerMy;
		interceptX = centerBx;
		interceptY = centerBy;
		*tiltOut = cylTilt;
		*phiOut = cylPhi;
	} else {
		// Use phi and tilt from the centers
		slopeX = centerMx;
		slopeY = centerMy;
		interceptX = centerBx;
		interceptY = centerBy;
		*tiltOut = centerTilt;
		*phiOut = centerPhi;
	}
	
	fprintf(stderr, "---- ---- built cylinder, writing into 3D image\n");
	for (i = 0; i < nSlices; i++){
		xCenter = slopeX * zLocs[i] + interceptX;
		yCenter = slopeY * zLocs[i] + interceptY;
		majorAxisLength = fabs(*rOut / cos(*tiltOut));
		sl = new_cylkSliceFromCylkImage (projectIm, 0, dim);
		ellipseSlice = new_cylkSliceLikeSliceWithFilledEllipse (&sl, majorAxisLength, *rOut, *phiOut, xCenter, yCenter);
		insertSliceIntoImage(&ellipseSlice, projectIm, i, dim);
	}
	
	// return phi and tilt off the real z and x axis...
	xCenter = centerMx * zLocs[nSlices/2] + centerBx;
	yCenter = centerMy * zLocs[nSlices/2] + centerBy;
	zCenter = zLocs[nSlices/2];
	
	switch (dim) {
		case 0:
			axis[0] = 1;
			axis[1] = centerMx;
			axis[2] = centerMy;
			anchor[0] = zCenter;
			anchor[1] = xCenter;
			anchor[2] = yCenter;
			break;
		case 1:
			axis[0] = centerMx;
			axis[1] = 1;
			axis[2] = centerMy;
			anchor[0] = xCenter;
			anchor[1] = zCenter;
			anchor[2] = yCenter;
			break;
		case 2:
			axis[0] = centerMx;
			axis[1] = centerMy;
			axis[2] = 1;
			anchor[0] = xCenter;
			anchor[1] = yCenter;
			anchor[2] = zCenter;
			break;
	}
	
	axisNorm = sqrt(pow(axis[0],2) + pow(axis[1],2) + pow(axis[2],2));
	for (i=0; i<3; i++) 
		axis[i] = axis[i]/axisNorm;
	
	*phiOut = atan2f(axis[1],axis[0]);
	*tiltOut = atan2f(sqrt(pow(axis[0],2)+pow(axis[1],2)),axis[2]);		
	return(0);
}

cylkSlice new_cylkSliceLikeSliceWithFilledEllipse(cylkSlice *like,
												  float majorAxis, 
												  float minorAxis, 
												  float phi,
												  float centerX, 
												  float centerY) {
	
	float f, f1[2], f2[2], rX, rY, r;
	int i;
	cylkSlice ellipseSlice = new_cylkSlice(like->nx, like->ny, like->dx, like->dy);
	
	//calculate the foci of the ellipse
	f = sqrt(pow(majorAxis, 2) - pow(minorAxis, 2));
	f1[0] = (f)  * (cos (phi)) + centerX;
	f1[1] = (f)  * (sin (phi)) + centerY;
	f2[0] = (-f) * (cos (phi)) + centerX;
	f2[1] = (-f) * (sin (phi)) + centerY;
	for (i=0; i<ellipseSlice.nVoxels; i++) {
		// fill in slice voxels inside ellipse
		rX = (i % ellipseSlice.nx) * ellipseSlice.dx;
		rY = (i / ellipseSlice.nx) * ellipseSlice.dy;
		r = sqrt(pow(rX-f1[0], 2) + pow(rY-f1[1], 2)) + sqrt(pow(rX-f2[0], 2) + pow(rY-f2[1], 2));
		ellipseSlice.pt[i] = r < (2 * majorAxis);
	}
	return(ellipseSlice);
}